package com.sz.admin.ai.util;

import com.mybatisflex.core.query.QueryWrapper;
import com.sz.admin.ai.factory.chatFactory.entity.MessageTypoeEnum;
import com.sz.admin.aichatmessage.pojo.po.AiChatMessage;
import com.sz.admin.aichatmessage.service.AiChatMessageService;
import lombok.AllArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * @描述:
 * @Author: TangYuan
 * @Date: 2025/3/5 15:06
 */
@Component
@AllArgsConstructor
public class MessageChatMemory implements ChatMemory {
    
    private final AiChatMessageService aiChatMessageService;
    
    public static Message toSpringAiMessage(AiChatMessage aiChatMessage) {
        MessageTypoeEnum anEnum = MessageTypoeEnum.getEnumByCapital(aiChatMessage.getType());
        if (anEnum == MessageTypoeEnum.SYSTEM) {
            return new SystemMessage(aiChatMessage.getContent());
        }
        if (anEnum == MessageTypoeEnum.ASSISTANT) {
            return new AssistantMessage(aiChatMessage.getContent());
        }
        if (anEnum == MessageTypoeEnum.USER) {
            return new UserMessage(aiChatMessage.getContent());
        }
        return null;
    }
    
    @Override
    public void add(String conversationId, Message message) {
        ChatMemory.super.add(conversationId, message);
    }
    
    @Override
    public void add(String conversationId, List<Message> messages) {
    
    }
    
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> list = new java.util.ArrayList<>(aiChatMessageService.list(QueryWrapper.create()
                        .eq("session_id", conversationId)
                        .ne("total_token" , 0)
                        .orderBy("id", false)
                        .limit(lastN))
                .stream()
                .map(MessageChatMemory::toSpringAiMessage)
                .filter(Objects::nonNull)
                .toList());
        
        Collections.reverse(list);
        
        System.out.println("Embedding-----------------------------------");
        list.forEach(message -> {
            System.out.println(message.getText());
        });
        System.out.println("Embedding-----------------------------------");
        
        
        return list;
    }
    
    @Override
    public void clear(String conversationId) {
        aiChatMessageService.remove(QueryWrapper.create().eq("session_id", conversationId));
    }
}
